package org.shj.algorithm.tree;

/**
 * 平衡二叉树，又叫 AVL 树。平衡二叉树的要求：对于每一个节点来说，其左右子树的
 * 高度差不超过 1。当插入或删除节点时，如果不满足要求，需要进行调整。
 * 平衡二叉树插入、查找、删除的时间复杂度的最好和最坏情况都是 O(logN)
 * @author Shen Huang Jian
 * @date 2019/5/21
 */
public class AVL {

    private Node<Integer> root;

    public void build(int[] arr){
        Node<Integer> val;
        for(int item : arr){
            val = new Node(item);
            val.setHeight(1);
            if(root == null){
                root = val;
            }else{
                insert(root, val);
            }
        }


    }

    public void insert(Node<Integer> node, Node<Integer> val){
        if(node.getVal() > val.getVal()){
            if(node.getLeftChild() == null){
                node.setLeftChild(val);
                val.setParent(node);

            }else{
                insert(node.getLeftChild(), val);
            }

            setNodeHeight(node);
            //如果左右子树的高度差大于2
            if(getHeightDiff(node) == 2){
                if(node.getLeftChild() != null && val.getVal() < node.getLeftChild().getVal()){
                    //说明插入新节点后是 左左 的情况
                    singleRotateLeft(node);
                }else{
                    doubleRotateLR(node); //左右
                }

            }

        }else{
            if(node.getRightChild() == null){
                node.setRightChild(val);
                val.setParent(node);

            }else{
                insert(node.getRightChild(), val);
            }

            setNodeHeight(node);
            //如果左右子树的高度差大于2
            if(getHeightDiff(node) == 2){
                if(node.getLeftChild() != null && val.getVal() < node.getLeftChild().getVal()){
                    //说明插入新节点后是 右左 的情况
                    doubleRotateRL(node);
                }else{
                    singleRotateRight(node); //右右
                }
            }
        }
    }

    public void middlePrint(){
        middle(root);
        System.out.println();
    }

    private void middle(Node<Integer> node){
        if(node == null){
            return;
        }
        middle(node.getLeftChild());
        System.out.print(node.getVal() + ", ");
        middle(node.getRightChild());
    }

    /**
     * 左左的情况，见图 AVL_左左.png 即图中节点6的左子树高于右子树，
     * 对于其左子树的左节点3，也是左子树高于右子树。
     * 此情况的旋转方法见图 AVL_左左_旋转.png
     * @param k2
     */
    private void singleRotateLeft(Node<Integer> k2){
        Node<Integer> k1 = k2.getLeftChild();
        Node<Integer> parent = k2.getParent();
        Node<Integer> k1Right = k1.getRightChild();

        //把 K2 做为 K1 的右孩子
        k1.setRightChild(k2);
        k2.setParent(k1);

        // 把 K1 的右孩子设为 K2 的左孩子
        k2.setLeftChild(k1Right);
        if(k1Right != null){
            k1Right.setParent(k2);
        }

        if(parent != null){
            if(parent.getLeftChild() == k2){
                parent.setLeftChild(k1);
            }else{
                parent.setRightChild(k1);
            }
        }else{
            root = k1;
        }
        k1.setParent(parent);

        setNodeHeight(k2);
        setNodeHeight(k1);
    }

    /**
     * 右右的情况，与 左左的情况 是对称的
     * @param k2
     */
    private void singleRotateRight(Node<Integer> k2){
        Node<Integer> k1 = k2.getRightChild();
        Node<Integer> parent = k2.getParent();
        Node<Integer> k1Left = k1.getLeftChild();

        //把 K2 做为 K1 的左孩子
        k1.setLeftChild(k2);
        k2.setParent(k1);

        // 把 K1 的左孩子设为 K2 的右孩子
        k2.setRightChild(k1Left);
        if(k1Left != null){
            k1Left.setParent(k2);
        }

        if(parent != null){
            if(parent.getLeftChild() == k2){
                parent.setLeftChild(k1);
            }else{
                parent.setRightChild(k1);
            }
        }else{
            root = k1;
        }
        k1.setParent(parent);

        setNodeHeight(k2);
        setNodeHeight(k1);
    }

    /**
     * 左右的情况，见图 AVL_左右.png
     * 旋转需要分两步，见图  AVL_左右_旋转.png
     *  第一步：转化成 左左 的情况，
     *  第二步：按左左的情况进行转换
     * @param k3
     */
    private void doubleRotateLR(Node<Integer> k3){
        singleRotateRight(k3.getLeftChild());
        singleRotateLeft(k3);
    }

    /**
     * 右左的情况，与 左右 的情况对称
     * @param k3
     */
    private void doubleRotateRL(Node<Integer> k3){
        singleRotateLeft(k3.getRightChild());
        singleRotateRight(k3);
    }

    private void setNodeHeight(Node<Integer> node){
        int leftH = node.getLeftChild() == null ? 0 : node.getLeftChild().getHeight();
        int rightH = node.getRightChild() == null ? 0 : node.getRightChild().getHeight();
        node.setHeight((leftH > rightH ? leftH : rightH) + 1);
    }

    private int getHeightDiff(Node<Integer> node){
        int left = node.getLeftChild() == null ? 0 : node.getLeftChild().getHeight();
        int right = node.getRightChild() == null ? 0 : node.getRightChild().getHeight();
        return (left > right) ? left - right : (right - left);
    }

    public static void main(String[] args){
        int[] arr = new int[]{35, 37, 47, 51, 58, 62, 73, 88, 93, 99};
        AVL avl = new AVL();
        avl.build(arr);

        avl.middlePrint();
    }
}
